package com.tianji.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.tianji.common.utils.UserContext;
import com.tianji.config.SystemPromptConfig;
import com.tianji.config.ToolResultHolder;
import com.tianji.constants.Constant;
import com.tianji.entity.ChatSession;
import com.tianji.enums.ChatEventTypeEnum;
import com.tianji.service.ChatService;
import com.tianji.service.ChatSessionService;
import com.tianji.vo.ChatEventVO;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.data.redis.core.BoundHashOperations;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static com.baomidou.mybatisplus.core.toolkit.Wrappers.lambdaQuery;
import static com.baomidou.mybatisplus.extension.toolkit.Db.updateById;

@Slf4j
@Service
@RequiredArgsConstructor
@ConditionalOnProperty(prefix = "tj.ai", name = "chat-type", havingValue = "ENHANCE")
public class ChatServiceImpl implements ChatService {

    private final ChatClient chatClient;

    private final SystemPromptConfig systemPromptConfig;

    private final ChatMemory chatMemory;

    private final ChatSessionService chatSessionService;

    // 存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全
    // 目前的版本暂时用Map实现，如果考虑分布式环境的话，可以考虑用redis来实现
    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();

    private final StringRedisTemplate  redisTemplate;

    private static final String GENERATE_STATUS_KEY = "GENERATE_STATUS_KEY";

    // 输出结束的标记
    private static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();

    // 向量数据库
    private final VectorStore vectorStore;




    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
        String conversationId = ChatService.getConversationId(sessionId);
        var hashOps = redisTemplate.boundHashOps(GENERATE_STATUS_KEY);
        StringBuilder outputBuilder = new StringBuilder();
        // 生成请求id
        var requestId = IdUtil.fastSimpleUUID();

        // 获取用户id
        var userId = UserContext.getUser();



        // 创建RAG增强
        var qaAdvisor = QuestionAnswerAdvisor.builder(this.vectorStore)
                .searchRequest(SearchRequest.builder()
                        .similarityThreshold(0.6d)
                        .topK(6)
                        .build()
                )
                .build();

        return chatClient.prompt()
                .system( s -> {
                    s.text(systemPromptConfig.getChatSystemMessage().get())
                            .param("now", DateUtil.now());
                })
                .advisors(advisor -> advisor
                        // 设置RAG增强
                        .advisors(qaAdvisor)
                        .param(chatMemory.CONVERSATION_ID, conversationId)
                )
                                                                   //传递用户id,因为tools的线程和服务线程不是一个线程
                                                                   //会导致鉴权问题
                .toolContext(Map.of(Constant.REQUEST_ID, requestId,Constant.USER_ID, userId))//通过工具上下文传递参数
                .user(question)
                .stream()
                .chatResponse()
                .doFirst( ()->{
                    //GENERATE_STATUS.put(sessionId, true);
                    hashOps.put(sessionId, "true");
                })
                .doOnError( throwable ->
                        //GENERATE_STATUS.remove(sessionId)
                        hashOps.delete(sessionId)
                )
                .doOnComplete(() ->
                        //GENERATE_STATUS.remove(sessionId)
                        hashOps.delete(sessionId)
                )
                // 保存停止输出的记录,防止会话记录保存不完整
                .doOnCancel(() -> saveStopHistoryRecord(conversationId,outputBuilder.toString()))
                .takeWhile(response -> {
                    return hashOps.get(sessionId) != null;
                })
                .map( chatResponse ->  {

                    // 对于响应结果进行处理，如果是最后一条数据，就把此次消息id放到内存中
                    // 主要用于存储消息数据，可以根据消息di获取的请求id，再通过请求id就可以获取到参数列表了
                    // 从而解决，在历史聊天记录中没有外参数的问题
                    var finishReason = chatResponse.getResult().getMetadata().getFinishReason();
                    if (StrUtil.equals(Constant.STOP, finishReason)) {
                        var messageId = chatResponse.getMetadata().getId();
                        ToolResultHolder.put(messageId, Constant.REQUEST_ID, requestId);
                    }

                    String text = chatResponse.getResult().getOutput().getText();
                    outputBuilder.append(text);
                    return ChatEventVO.builder()
                            .eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                .concatWith(Flux.defer(() -> {
                    // 通过请求id获取到参数列表，如果不为空，就将其追加到返回结果中
                    var map = ToolResultHolder.get(requestId);
                    if (CollUtil.isNotEmpty(map)) {
                        ToolResultHolder.remove(requestId); // 清除参数列表

                        // 响应给前端的参数数据
                        var chatEventVO = ChatEventVO.builder()
                                .eventData(map)
                                .eventType(ChatEventTypeEnum.PARAM.getValue())
                                .build();
                        return Flux.just(chatEventVO, STOP_EVENT);
                    }
                    return Flux.just(STOP_EVENT);
                }).doFinally(
                        signalType -> {
                            // 异步更新对话信息 需要调整chatSessionService.update 接口
                            chatSessionService.update(sessionId,question,userId,outputBuilder.toString());
                        }
                ));

    }

    private void saveStopHistoryRecord(String conversationId, String content) {
        chatMemory.add(conversationId, new AssistantMessage(content));
    }






    @Override
    public void stop(String sessionId) {
        redisTemplate.boundHashOps(GENERATE_STATUS_KEY).rename(sessionId);
    }


}